# -*- coding: utf-8 -*-
# @Time : 2020/11/1 13:56
# @Author : wudeyang
# @email :wudeyang@sjtu.edu.cn
# @Description:


from torch.utils.data import DataLoader
from torchvision import transforms
import copy
from utils import  show_img
from Data_Loader import dataset
import matplotlib.pyplot as plt
import pathlib
import sys
import os

def GET_DATA_LOADER(moudle_args):
    all_supported_dataset=['Synth80k','Synth3D']

    if moudle_args['type'] not in all_supported_dataset:

        sys.exit('not supported dataset')
    if moudle_args['type'] =='Synth80k':
        return get_dataloader_Synth80k(moudle_args)
    elif moudle_args['type']=='Synth3D':
        return get_dataloader_Synth3D(moudle_args)
    else:
        sys.exit('unknow bug')



def get_dataloader_Synth80k( module_args):
    # 合成数据 Synth80k
    train_transfroms = transforms.Compose([
        transforms.ColorJitter(brightness=0.5),
        transforms.ToTensor()
    ])
    # 创建数据集\
    dataset_args = copy.deepcopy(module_args)

    data_path = dataset_args['data_path']

    train_dataset=dataset.Synth80k(data_path,moudle_args['input_size'],train_transfroms)

    train_loader = DataLoader(dataset=train_dataset, batch_size=dataset_args['train_batch_size'],
                              shuffle=dataset_args['shuffle'], num_workers=dataset_args['num_workers'],pin_memory=True)
    train_loader.dataset_len=len(train_dataset)


    return train_loader


def get_dataloader_Synth3D(moudle_args):
    train_transforms = transforms.Compose([
        transforms.ColorJitter(brightness=0.5),
        transforms.ToTensor()
    ])
    # 创建数据集\
    dataset_args = copy.deepcopy(moudle_args)

    train_data_path = dataset_args['data_path'] # 要求是一个txt文件，文件格式如下
    # train_data 数据格式
    # [('/data/Synth3D-10K/img/5545.jpg', '/data/Synth3D-10K/label/5545.txt'),
    #  ('/data/Synth3D-10K/img/2395.jpg', '/data/Synth3D-10K/label/2395.txt'),
    #  ('/data/Synth3D-10K/img/8809.jpg', '/data/Synth3D-10K/label/8809.txt')]
    train_data = []

    with open(train_data_path, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            line = line.strip('\n').replace('.jpg ', '.jpg\t').split('\t')
            if len(line) > 1:
                img_path = pathlib.Path(line[0].strip(' '))
                label_path = pathlib.Path(line[1].strip(' '))
                if img_path.exists() and img_path.stat().st_size > 0 and label_path.exists() and label_path.stat().st_size > 0:
                    train_data.append((str(img_path), str(label_path)))

    train_dataset=dataset.Synth3D(train_data,input_size=dataset_args['input_size'],transform=train_transforms)
    train_loader = DataLoader(dataset=train_dataset, batch_size=dataset_args['train_batch_size'],
                              shuffle=dataset_args['shuffle'], num_workers=dataset_args['num_workers'],pin_memory=True)
    train_loader.dataset_len=len(train_dataset)


    return train_loader


if __name__=='__main__':
    os.environ['CUDA_VISIBLE_DEVICES'] = '3'
    moudle_args={
            "type": "Synth80k",
            "data_path": "/data/SynthText",
            "train_batch_size": 1,
            "shuffle": False,
            "pin_memory": False,
            "num_workers": 12,
            "input_size":[720,720]
        }


    train_loader =GET_DATA_LOADER (moudle_args)
    for i, (img, mask, gauss) in enumerate(train_loader):
        if i>0:
            break

        print(img.shape)

        print(mask.shape)
        print(gauss.shape)
        show_img(img[0].permute(1, 2, 0), color=True)
        show_img(mask[0])
        show_img(gauss[0])
        plt.show()


